package grpc

import (
	"context"
	"net"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/reflection"
)

// Server is a gRPC server wrapper.
type Server struct {
	*grpc.Server
	log     grpclog.LoggerV2
	network string
	address string
	timeout time.Duration
	health  *health.Server
}

// Start the gRPC server.
func (s *Server) Start(ctx context.Context) error {
	lis, err := net.Listen(s.network, s.address)
	if err != nil {
		return err
	}
	s.log.Infof("[gRPC] server listening on: %s", lis.Addr().String())
	s.health.Resume()
	return s.Serve(lis)
}

// Stop the gRPC server.
func (s *Server) Stop(ctx context.Context) error {
	s.health.Shutdown()
	s.GracefulStop()
	s.log.Info("[gRPC] server stopping")
	return nil
}

func NewServer(opts ...ServerOption) *Server {
	srv := &Server{
		network: "tcp",
		address: ":0",
		timeout: time.Second * 3,
		health:  health.NewServer(),
	}
	for _, o := range opts {
		o(srv)
	}
	srv.Server = grpc.NewServer(grpc.UnaryInterceptor(unaryServerInterceptor(srv)))
	grpc_health_v1.RegisterHealthServer(srv.Server, srv.health)
	reflection.Register(srv.Server)
	return srv
}

func unaryServerInterceptor(s *Server) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		var cancel context.CancelFunc
		if s.timeout > 0 {
			ctx, cancel = context.WithTimeout(ctx, s.timeout)
		} else {
			ctx, cancel = context.WithCancel(ctx)
		}
		defer cancel()
		return handler(ctx, req)
	}
}
